import sqlite3
import mysql.connector
import psycopg2
import pandas as pd
from sqlalchemy import create_engine
from flask import g
from config import Config

# 数据库连接缓存
DATABASE = Config.SQLALCHEMY_DATABASE_URI.replace('sqlite:///', '')

def get_db():
    """获取本地数据库连接"""
    db = getattr(g, '_database', None)
    if db is None:
        db = g._database = sqlite3.connect(DATABASE)
        db.row_factory = sqlite3.Row
    return db

def close_db(exception):
    """关闭数据库连接"""
    db = getattr(g, '_database', None)
    if db is not None:
        db.close()

def init_db(app):
    """初始化数据库"""
    with app.app_context():
        db = get_db()
        with app.open_resource('schema.sql', mode='r') as f:
            db.cursor().executescript(f.read())
        db.commit()
    
    # 注册关闭数据库的函数
    app.teardown_appcontext(close_db)

class DatabaseConnector:
    """数据库连接器，支持多种数据库类型"""
    
    def __init__(self, connection_params):
        self.connection_params = connection_params
        self.connection = None
        self.db_type = connection_params.get('db_type', 'mysql')
    
    def connect(self):
        """建立数据库连接"""
        try:
            if self.db_type == 'mysql':
                self.connection = mysql.connector.connect(
                    host=self.connection_params['host'],
                    port=self.connection_params.get('port', 3306),
                    user=self.connection_params['username'],
                    password=self.connection_params['password'],
                    database=self.connection_params['database']
                )
            elif self.db_type == 'postgresql':
                self.connection = psycopg2.connect(
                    host=self.connection_params['host'],
                    port=self.connection_params.get('port', 5432),
                    user=self.connection_params['username'],
                    password=self.connection_params['password'],
                    dbname=self.connection_params['database']
                )
            elif self.db_type == 'sqlserver':
                conn_str = (
                    f"mssql+pymssql://{self.connection_params['username']}:{self.connection_params['password']}@"
                    f"{self.connection_params['host']}:{self.connection_params.get('port', 1433)}/"
                    f"{self.connection_params['database']}"
                )
                engine = create_engine(conn_str)
                self.connection = engine.connect()
            elif self.db_type == 'sqlite':
                self.connection = sqlite3.connect(self.connection_params['database'])
            else:
                raise ValueError(f"不支持的数据库类型: {self.db_type}")
            
            return True
        except Exception as e:
            print(f"数据库连接失败: {str(e)}")
            return False
    
    def execute_query(self, query, params=None):
        """执行查询并返回结果"""
        if not self.connection:
            if not self.connect():
                return None, "连接数据库失败"
        
        try:
            cursor = self.connection.cursor()
            cursor.execute(query, params or ())
            
            # 如果是SELECT语句，返回结果
            if query.strip().upper().startswith('SELECT'):
                columns = [column[0] for column in cursor.description]
                data = cursor.fetchall()
                result = [dict(zip(columns, row)) for row in data]
                return result, None
            else:
                # 对于INSERT, UPDATE, DELETE等语句，提交事务
                self.connection.commit()
                return {"rows_affected": cursor.rowcount}, None
        except Exception as e:
            self.connection.rollback()
            return None, str(e)
        finally:
            cursor.close()
    
    def close(self):
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()
            self.connection = None

def allowed_file(filename):
    """检查文件是否为允许的类型"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in Config.ALLOWED_EXTENSIONS

def read_file_to_dataframe(file_path):
    """读取文件内容到DataFrame"""
    try:
        if file_path.endswith('.csv'):
            return pd.read_csv(file_path)
        elif file_path.endswith(('.xlsx', '.xls')):
            return pd.read_excel(file_path)
        elif file_path.endswith('.json'):
            return pd.read_json(file_path)
        else:
            return None
    except Exception as e:
        print(f"读取文件失败: {str(e)}")
        return None
